from csv import excel

import pandas as pd
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import seaborn as sns


def perform_kmeans_clustering(file_path):
    # Set plot style
    sns.set(style="whitegrid")


    # 1. Load the data
    df = pd.read_excel(file_path)

    # 2. Pre-process data
    t90_data = df[['t90']].dropna()
    X = t90_data.values.reshape(-1, 1)

    # 3. Find optimal K using the elbow method
    inertia = []
    K_range = range(2, 11)
    for k in K_range:
        kmeans = KMeans(n_clusters=k, init='k-means++', random_state=42, n_init=10)
        kmeans.fit(X)
        inertia.append(kmeans.inertia_)


    # We choose K=4 based on the elbow curve
    k = 5
    kmeans = KMeans(n_clusters=k, init='k-means++', random_state=42, n_init=10)
    df['cluster'] = kmeans.fit_predict(X)

    # 5. Analyze results and output BMI ranges
    bmi_ranges = df.groupby('cluster')['基准BMI'].agg(['min', 'max']).reset_index()
    bmi_ranges.columns = ['分组', 'BMI_min', 'BMI_max']
    print("每个分组的BMI范围:")
    print(bmi_ranges)

    # Save clustered data
    df.to_csv('data_with_clusters.csv', index=False)
    cluster_centers = kmeans.cluster_centers_
    for i, center in enumerate(sorted(cluster_centers.flatten())):
        print(f"聚类 {i}: {center:.4f}")
    # 6. Visualize results
    # Scatter plot of t90 vs. BMI, colored by cluster
    plt.figure(figsize=(12, 8))
    sns.scatterplot(data=df, x='t90', y='基准BMI', hue='cluster', palette='viridis', s=100, alpha=0.7)
    plt.title('t90 vs. 基准BMI (按聚类分组)')
    plt.xlabel('t90')
    plt.ylabel('基准BMI')
    plt.legend(title='分组')
    plt.savefig('t90_bmi_scatterplot.png')
    plt.close()

    # Box plot of BMI for each cluster
    plt.figure(figsize=(12, 8))
    sns.boxplot(data=df, x='cluster', y='基准BMI', palette='viridis')
    plt.title('每个分组的基准BMI分布')
    plt.xlabel('分组')
    plt.ylabel('基准BMI')
    plt.savefig('bmi_boxplot_by_cluster.png')
    plt.close()

if __name__ == '__main__':
    # Assuming 'data_after_q2_2.xlsx - Sheet1.csv' is in the same directory
    perform_kmeans_clustering('data_after_q2_2_noisy_2.xlsx')
